#include "dyn_planarBilliard.h"

PlanarBilliard::PlanarBilliard(Functor& _boundary,
                               Functor& _diffBoundary,
                               Functor& _diff2Boundary,
                               Quad &_quadrature,
                               colvec _initCond,
                               double _tolQuad,
                               double _tolBis,
                               double _tolNew,
                               int _numberBounces)
    : boundary(_boundary), diffBoundary(_diffBoundary), diff2Boundary(_diff2Boundary),
      quadrature(_quadrature), arcLength(_boundary,_diffBoundary)
{
    // Set the values for this billiard object.
    setInitialConditions(_initCond);
    setToleranceQuadrature(_tolQuad);
    setToleranceBisection(_tolBis);
    setToleranceNewton(_tolNew);
    setNumberBounces(_numberBounces);

    // Sets the variables of the NextImpact object
    nextImpact = new NextImpactFunction(_boundary);

    // Sets the variables of the NextDiffImpact object
    nextImpactDiff = new NextImpactDiff(_boundary,_diffBoundary);

    // Sets the variables of the Hunt object.
    hunt = new Hunting(*nextImpact,*nextImpactDiff,0.0,0.0,0.0,0.0,0.0,0.0,0);

    // We compute the total arc length and store it.
    quadrature.setFunctor(arcLength);
    quadrature.setA(0.0);
    quadrature.setB(2.0*datum::pi);
    quadrature.setTol(tolQuad);
    totalArc = quadrature.doQuadrature();
}

double PlanarBilliard::computeArcLength(double phi)
{
    quadrature.setA(0.0);
    quadrature.setB(user_fmod(phi,2.*datum::pi));

    return quadrature.doQuadrature()/totalArc;
}

double PlanarBilliard::computeMomentum(double alpha)
{
    return cos(alpha);
}

colvec PlanarBilliard::computeNextImpact(colvec previousImpact)
{
    // This algorithm is based on results by Korsch.
    // We define variables containing the previous angles.
    double phi_n = previousImpact(0);
    double alpha_n = previousImpact(1);

    // We compute the angle between the positive direction of the
    // tangent and the radial ray.
    double theta_n = datum::pi
            -acos(diffBoundary(phi_n)*pow(pow(boundary(phi_n),2.0)
                                          +pow(diffBoundary(phi_n),2.0),-0.5));

    // The direction of the trajectory is given by (from phi=0):
    double beta_n = datum::pi+phi_n+alpha_n-theta_n;

    // We can now have the equation of the trajectory as:
    nextImpact->setDirection(beta_n);
    nextImpact->setInputImpactPoint(phi_n);

    // And its derivative.
    nextImpactDiff->setDirection(beta_n);
    nextImpactDiff->setImpactPoint(phi_n);

    // Which we solve with a combination of bisection and Newton-Raphson.
    hunt->setInitialValue(phi_n+1.0e-15);
    hunt->setEndPoint(phi_n+2.0*datum::pi);
    hunt->setScanningSpeed(1.0e-1);
    hunt->setToleranceBis(tolBis);
    hunt->setTolerance(tolNew);
    hunt->setMaxValue(1.0);
    hunt->setMaxIterations(1e5);

    int test = hunt->solve();

    while (test)
    {
        hunt->setScanningSpeed(hunt->getScanningSpeed()/2.0);
        test = hunt->solve();

        if (hunt->getScanningSpeed()<1.0e-15)
            break;
    }

    if (test)
        cout << "Root-finding failed... " << endl;

    double phi_n1 = user_fmod(hunt->getRoot(),2.*datum::pi);

    // We compute the new angle from the positive tangent and the ray.
    double theta_n1 = datum::pi
            -acos(diffBoundary(phi_n1)*pow(pow(boundary(phi_n1),2.0)
                                           +pow(diffBoundary(phi_n1),2.0),-0.5));

    // And the new alpha, and the work is done.
    double alpha_n1 = phi_n1-phi_n+theta_n-theta_n1-alpha_n;

    colvec results(2);
    results(0) = phi_n1;
    results(1) = user_fmod(alpha_n1,datum::pi);

    return results;
}

mat PlanarBilliard::computeTrajectory()
{
    // From an initial condition, we compute
    // the whole trajectory of a ball for a
    // given number of bounces.
    mat results(numberBounces+1,4);

    // Initial data.
    results(0,0) = initCond(0);
    results(0,1) = initCond(1);
    results(0,2) = computeArcLength(initCond(0));
    results(0,3) = computeMomentum(initCond(1));

    for (unsigned int i=1;i<results.n_rows;i++)
    {
        // We compute the next step.
        colvec oldImpact = results.submat(span(i-1),span(0,1)).t();
        colvec newImpact = computeNextImpact(oldImpact);

        results.submat(span(i),span(0,1)) = newImpact.t();
        results(i,2) = computeArcLength(results(i,0));
        results(i,3) = computeMomentum(results(i,1));
    }

    return results;
}

mat PlanarBilliard::huntUPOs(double lambdaInit,
                             double tolSTA,
                             int period,
                             int maxIterations)
{
    // Preparation phase: we generate the reflection matrices necessary.
    cube reflectionMatrices = prepareReflectionMatrices();

    // We prepare a matrix to hold the fixed points.
    mat fixedPoints(period+1,2);
    fixedPoints.row(0) = getInitialConditions().t();

    /* We now use the standard STA algorithm, which
     * goes as follows. Using the potential fixed point
     * found by recurrence (this is unlikely, since
     * billiard are in essence Hamiltonian systems;
     * it might take an infinitely long time to return
     * to the same point), we iterate the STA system
     * until we converge to a fixed point. If the corrections
     * are, after some initial transient, still high, we
     * use another reflection matrix. After having used all the
     * matrices, we halve the stepsize lambda. If this still doesn't work,
     * then we quit the program.
     *
     * After a fixed point is found, we iterate the point once to find
     * the next initial condition. */
    for (int i=0;i<period;i++)
    {
        // We assign a temporary variable as initial condition.
        colvec oldStep = fixedPoints.row(i).t();

        // Some bookkeeping first
        double lambda = lambdaInit;
        unsigned int slice (0),iterations(0);
        // We iterate the STA system until convergence / divergence.
        while(1)
        {
            // Using the initial condition (and subsequent values),
            // we compute the STA correction.
            colvec dx = computeSTAcorrection(lambda,reflectionMatrices.slice(slice),oldStep,period);

            // And apply it to the initial condition.
            oldStep += dx;
            oldStep(0) = user_fmod(oldStep(0),2.*datum::pi);
            oldStep(1) = user_fmod(oldStep(1),datum::pi);

            // Do something if the correction is too high after
            // an initial transient.
            if (iterations>maxIterations || dx(0)!=dx(0))
            {
                // We change matrix and begin anew.
                slice++;
                oldStep = fixedPoints.row(i).t();
                iterations = 0;
            }

            // If we have tried all matrices, half stepsize
            // and try again.
            if (slice==reflectionMatrices.n_slices)
            {
                // Half stepsize and begin anew.
                lambda /= 2.0;
                slice=0;

                oldStep = fixedPoints.row(i).t();
            }

            // If the stepsize is smaller than the tolerance,
            // we have a problem.
            if (lambda < 1.0e-15)
            {
                cout << "Failed to converge to a fixed point." << endl;
                break;
            }

            // Check if we have converged.
            if (norm(dx,2) < tolSTA)
            {
                // We have converged! Store the fixed point.
                fixedPoints.row(i) = oldStep.t();

                // Iterate once the next point in the orbit.
                fixedPoints.row(i+1) = computeNextImpact(fixedPoints.row(i).t()).t();
                break;
            }

            // Some more bookkeeping
            iterations++;
        }
    }

    // Once we have the first point, we iterate them
    for (int i=1;i<=period;i++)
    {
        fixedPoints.row(i) = computeNextImpact(fixedPoints.row(i-1).t()).t();
    }
    // Once we have the whole matrix, we compute the
    // associated s,p variables.
    mat orbit(period+1,4);

    for (int i=0;i<=period;i++)
    {
        orbit(i,0) = user_fmod(fixedPoints(i,0),2.0*datum::pi);
        orbit(i,1) = user_fmod(fixedPoints(i,1),datum::pi);
        orbit(i,2) = computeArcLength(orbit(i,0));
        orbit(i,3) = computeMomentum(orbit(i,1));
    }

    return orbit;
}

cube PlanarBilliard::prepareReflectionMatrices()
{
    cube reflectionMatrices(2,2,5);
    mat reflection1(2,2), reflection2(2,2),reflection3(2,2),
            reflection4(2,2),reflection5(2,2);
    reflection1=eye(2,2);
    reflection2 << -1.0 << 0.0 << endr
                << 0.0  << 1.0 << endr;

    reflection3 << 1.0 << 0.0 << endr
                << 0.0 << -1.0 << endr;

    reflection4 << 0.0 << -1.0 << endr
                << -1.0 << 0.0 << endr;

    reflection5 << 0.0 << 1.0 << endr
                << 1.0 << 0.0 << endr;

    reflectionMatrices.slice(0) = reflection1;
    reflectionMatrices.slice(1) = reflection2;
    reflectionMatrices.slice(2) = reflection3;
    reflectionMatrices.slice(3) = reflection4;
    reflectionMatrices.slice(4) = reflection5;

    return reflectionMatrices;
}

colvec PlanarBilliard::computeSTAcorrection(double lambda,
                                            mat reflection,
                                            colvec previousStep,
                                            int period)
{
    // We compute the next impact of the billiard.
    colvec nextStep = previousStep;
    for (int i=0;i<period;i++)
    {
          nextStep = computeNextImpact(nextStep);
    }

    colvec correction = lambda*reflection*(nextStep-previousStep);

    return correction;
}

double PlanarBilliard::computeRadiusOfCurvature(double phi)
{
    double numerator = pow(boundary(phi),2.0)+pow(diffBoundary(phi),2.0);
    double denominat = pow(boundary(phi),2.0)+2.0*pow(diffBoundary(phi),2.0)-boundary(phi)*diff2Boundary(phi);

    return pow(numerator,1.5)*pow(denominat,-1.0);
}

mat PlanarBilliard::computeJacobianMatrix(colvec impact)
{
    // We compute the next impact.
    colvec newImpact = computeNextImpact(impact);

    // We define the individual variables.
    double phi_n = impact(0);
    double alpha_n = impact(1);
    double phi_n1 = newImpact(0);
    double alpha_n1 = newImpact(1);
    double rho0 = computeRadiusOfCurvature(phi_n);
    double rho1 = computeRadiusOfCurvature(phi_n1);

    // We compute the chord length.
    double l01 = sqrt(pow(boundary(phi_n),2.0)+pow(boundary(phi_n1),2.0)
                      -2.0*boundary(phi_n)*boundary(phi_n1)*cos(phi_n1-phi_n));

    // We compute the matrix elements.
    double a11 = -sin(alpha_n)/sin(alpha_n1)+l01/(sin(alpha_n1)*rho0);
    double a12 = -l01/(sin(alpha_n)*sin(alpha_n1)*totalArc);
    double a21 = totalArc*(-l01/(rho0*rho1)+sin(alpha_n1)/rho0+sin(alpha_n)/rho1);
    double a22 = -sin(alpha_n1)/sin(alpha_n)+l01/(sin(alpha_n)*rho1);

    // And build the matrix.
    mat jacobian(2,2);

    jacobian << a11 << a12 << endr
             << a21 << a22 << endr;

    return jacobian;
}
